from ultralytics import YOLO
import cv2

# 加载训练好的模型
model = YOLO(model='D:/Project/All_Project/Python/blct/runs/train/cs8/weights/best.pt')

# 推理函数
def inference_and_display(image_path):
    # 读取图像
    frame = cv2.imread(image_path)
    if frame is None:
        print(f"Error: Unable to load image at {image_path}")
        return

    # 使用模型进行推理
    results = model(frame)  # 推理

    # 解析推理结果
    for result in results:
        # 绘制检测框和标签
        boxes = result.boxes.xyxy.cpu().numpy()  # 检测框坐标 (x1, y1, x2, y2)
        confidences = result.boxes.conf.cpu().numpy()  # 置信度
        class_ids = result.boxes.cls.cpu().numpy()  # 类别ID

        for box, confidence, class_id in zip(boxes, confidences, class_ids):
            x1, y1, x2, y2 = map(int, box)  # 转换为整数
            label = f"{model.names[int(class_id)]} {confidence:.2f}"  # 标签和置信度

            # 绘制检测框
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            # 绘制标签
            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

    # 显示结果
    cv2.imshow("YOLO Inference", frame)
    cv2.waitKey(0)  # 等待按键
    cv2.destroyAllWindows()  # 关闭窗口

# 测试推理
if __name__ == '__main__':
    image_path = "D:/Project/All_Project/Python/blct/YOLOv11/datasets/cs/images/3.png"  # 替换为你的测试图像路径
    inference_and_display(image_path)